package main

import (
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/huawen0327/golee"
)

var (
	upgrader = websocket.Upgrader{
		// 读取存储空间大小
		ReadBufferSize: 65536,
		// 写入存储空间大小
		WriteBufferSize: 65536,
		// 允许跨域
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	// DevicesMap 设备指针
	DevicesMap = DevicesStruct{
		Devices: make(map[string]*websocket.Conn),
	}
	// Upgrade header
	Upgrade  = []string{"websocket"}
	data     []byte
	pingData = []byte("aaaabbbbcccc")
)

// Connetion websocket的连接
type Connetion struct {
	wbsCon *websocket.Conn
	mutex  sync.Mutex
}

// DevicesStruct 存储链接的map
type DevicesStruct struct {
	Devices map[string]*websocket.Conn
	Mutex   sync.Mutex
}

// 定时ping
func (conn *Connetion) ping(ticker *time.Ticker, key string) {
	for {
		<-ticker.C
		fmt.Println("ping")
		conn.mutex.Lock()
		if err := conn.wbsCon.WriteMessage(websocket.PingMessage, pingData); err != nil {
			goto ERR // 发送消息失败，关闭连接
		}
		fmt.Println("pong")
		conn.mutex.Unlock()
	}
ERR:
	// 关闭连接
	if err := conn.wbsCon.Close(); err != nil {
		fmt.Println("websocket had closed")
	}
	conn.mutex.Unlock()
	ticker.Stop()

	DevicesMap.Mutex.Lock()
	delete(DevicesMap.Devices, key)
	DevicesMap.Mutex.Unlock()
}

func (conn *Connetion) recv(key string, wg *sync.WaitGroup) {
	var data []byte
	var err error
	for {
		conn.mutex.Lock()
		if _, data, err = conn.wbsCon.ReadMessage(); err != nil {
			goto ERR // 跳转到关闭连接
		}
		if err = conn.wbsCon.WriteMessage(websocket.TextMessage, data); err != nil {
			goto ERR // 发送消息失败，关闭连接
		}
		conn.mutex.Unlock()
	}
ERR:
	// 关闭连接
	if err = conn.wbsCon.Close(); err != nil {
		fmt.Println("websocket had closed")
	}
	conn.mutex.Unlock()
	DevicesMap.Mutex.Lock()
	delete(DevicesMap.Devices, key)
	DevicesMap.Mutex.Unlock()
	wg.Done()
}

var connWS = func(c *golee.Context) {
	fmt.Println(c.Req.Header["Upgrade"])
	for key, value := range c.Req.Header["Upgrade"] {
		if value != Upgrade[key] {
			c.StringResponse(404, "404 NOT FOUND: %s\n", c.Path) // 获取连接失败直接返回
			return
		}
	}
	var wbsCon Connetion
	if wbsCon.wbsCon, err = upgrader.Upgrade(c.Writer, c.Req, nil); err != nil {
		c.StringResponse(404, "404 NOT FOUND: %s\n", c.Path) // 获取连接失败直接返回
		return
	}
	key := c.GET("deviceid")
	if key == "" {
		// 关闭连接
		if err := wbsCon.wbsCon.Close(); err != nil {
			fmt.Println("websocket had closed")
		}
		// c.StringResponse(404, "404 NOT FOUND: %s\n", c.Path) // 获取连接失败直接返回
		return
	}

	DevicesMap.Mutex.Lock()
	DevicesMap.Devices[key] = wbsCon.wbsCon
	DevicesMap.Mutex.Unlock()

	var wg sync.WaitGroup
	// cancel := make(chan bool)

	// ticker := time.NewTicker(time.Second * 10)
	// go wbsCon.ping(ticker, key)

	wg.Add(1)
	go wbsCon.recv(key, &wg)
	wg.Wait()
	// close(cancel)

	// 关闭连接
	fmt.Println("close")
	// 关闭连接
	if err := wbsCon.wbsCon.Close(); err != nil {
		fmt.Println("websocket had closed")
	}

	DevicesMap.Mutex.Lock()
	delete(DevicesMap.Devices, key)
	DevicesMap.Mutex.Unlock()
}
